#!/usr/bin/env python
# coding=utf-8
# 从NCBI的refseq注释中 过滤获得最长转录本
# 注意 这里识别一个 gene 为蛋白编码基因是必须具有cds的gff行


import argparse
import sys
import re


parser = argparse.ArgumentParser(
    description='''从refseq给的 gff3中 提取最长蛋白序列 的基因 的mRNA 将其转成一个gff3文件 用于后续的同源聚类分析
    用法:
    fetch_longest_pep_from_refseq.py -i GCF_002204515.2_AaegL5.0_genomic.gff -o longest_isoform.gff3
    由大天才于2021年7月10日创建于浙江农业大学''')

parser.add_argument('-i',
                help='必须给定，输入的refseq gff3文件 ')


parser.add_argument('-o',
                help='必须给定，输出文件')



args = parser.parse_args()

if not args.i or not args.o:
    parser.print_help()
    sys.exit()



infile = args.i

outfile= args.o






lista = []

gff_genes = {}

mRNA_ID_to_gene_ID = {}

item_property_pattern = re.compile(r'([^=|^;]+)=([^;]+)')

with open(infile) as fila:
    for i in fila:
        k = i.strip().split('#')[0].split('\t')


        if k[0] != '' and len(k)>=8 and k[3].isdigit() and k[4].isdigit():
            # 读取id
          
            item_dic = {i[0]:i[1] for i in item_property_pattern.findall(k[8])}

            # 如果是基因 就建立该基因的 mRNA字典

            if k[2] =='gene':
                
                gff_genes[item_dic['ID']] = {
                'self': k,
                'mRNAs':{},
                }

            if k[2] == 'mRNA':
                if item_dic['Parent'] in gff_genes:
                    pass

                else:
                    # 有时 会没写 
                    gff_genes[item_dic['Parent']] = {
                    'self': [],
                    'mRNAs':{},
                    }

                gff_genes[item_dic['Parent']]['mRNAs'][item_dic['ID']] = {
                'self': k,
                'length':0,
                'cds':[],
                'exon':[],
                }

                mRNA_ID_to_gene_ID[item_dic['ID']] = item_dic['Parent']
            if k[2] == 'exon':

                if item_dic['Parent'] not in mRNA_ID_to_gene_ID:
                    continue

                gff_genes[mRNA_ID_to_gene_ID[item_dic['Parent']]]['mRNAs'][item_dic['Parent']]['exon'].append(k)
            
            if k[2].upper() == 'CDS':

                if item_dic['Parent'] not in mRNA_ID_to_gene_ID:
                    continue

                length = abs(int(k[3])-int(k[4]))

                gff_genes[mRNA_ID_to_gene_ID[item_dic['Parent']]]['mRNAs'][item_dic['Parent']]['cds'].append(k)
                gff_genes[mRNA_ID_to_gene_ID[item_dic['Parent']]]['mRNAs'][item_dic['Parent']]['length'] += length



# 写出最长的那个mRNA


outhandel = open(outfile,'w')
outhandel.write('##gff-version 3\n')

for i in gff_genes:

    longest_mRNA_id = ''
    max_length = 0

    
    for j in gff_genes[i]['mRNAs']:


        if gff_genes[i]['mRNAs'][j]['length'] > max_length:

            max_length = gff_genes[i]['mRNAs'][j]['length']

            longest_mRNA_id = j



    if max_length == 0:

        continue

    if gff_genes[i]['self']!=[]:
        outhandel.write('\t'.join(gff_genes[i]['self'])+'\n')
    outhandel.write('\t'.join(gff_genes[i]['mRNAs'][longest_mRNA_id]['self'])+'\n')
    for k in gff_genes[i]['mRNAs'][longest_mRNA_id]['exon']:
        outhandel.write('\t'.join(k)+'\n')
    for k in gff_genes[i]['mRNAs'][longest_mRNA_id]['cds']:
        outhandel.write('\t'.join(k)+'\n')

outhandel.close()
